package org.dromara.common.websocket.handler;

import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.dromara.common.core.domain.model.LoginUser;
import org.dromara.common.websocket.holder.WebSocketSessionHolder;
import org.dromara.common.websocket.utils.UUIDToByteArray;
import org.dromara.common.websocket.utils.WebSocketUtils;
import org.dromara.llm.domain.LlmSender;
import org.dromara.llm.domain.replyDTO;
import org.dromara.llm.service.IAiChatService;
import org.dromara.llm.service.impl.AITranslateServiceImpl;
import org.dromara.llm.util.AudioUtils;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import javax.annotation.Resource;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.*;

import static org.dromara.common.websocket.constant.WebSocketConstants.LOGIN_USER_KEY;
import static org.dromara.common.websocket.constant.WebSocketConstants.PEER_ID;


/**
 * WebSocketHandler 实现类
 *
 * @author zendwang
 */
@Slf4j
public class PlusWebSocketHandler extends AbstractWebSocketHandler {
    @Resource
    private IAiChatService aiChatService;

    @Resource
    private AITranslateServiceImpl aiTranslateService;

    // 核心线程数：根据CPU核心数和IO密集度调整
    private static final int CORE_POOL_SIZE = Runtime.getRuntime().availableProcessors() * 2;
    // 最大线程数：允许临时扩展的线程数量
    private static final int MAX_POOL_SIZE = CORE_POOL_SIZE * 2;
    // 空闲线程存活时间
    private static final long KEEP_ALIVE_TIME = 60L;
    // 任务队列容量：避免使用无界队列防止OOM
    private static final int QUEUE_CAPACITY = 1000;

    // 创建自定义线程池
    private static final ExecutorService WORKER_POOL = new ThreadPoolExecutor(
        CORE_POOL_SIZE,
        MAX_POOL_SIZE,
        KEEP_ALIVE_TIME,
        TimeUnit.SECONDS,
        new LinkedBlockingQueue<>(QUEUE_CAPACITY),
        new ThreadFactoryBuilder()
            .setNameFormat("websocket-worker-%d")
            .setDaemon(true)
            .build(),
        new ThreadPoolExecutor.CallerRunsPolicy()
    );

    //视频通话连接记录
    private static final Map<String, WebSocketSession> peerIdSessionMap = new ConcurrentHashMap<>();

    //视频通话列表
    private static final Map<String, String> videoAiMap = new ConcurrentHashMap<>();

    //语言
    private static final Map<String, String> languageMap = new ConcurrentHashMap<>();

    //语言
    private static final Map<String, Boolean> statusMap = new ConcurrentHashMap<>();


    /**
     * 连接成功后
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws IOException {
        //视频通话
        String peerId = (String) session.getAttributes().get(PEER_ID);
        if (peerId != null) {
            log.info("on open:the session is is :{},the peer id is:{}", session.getId(), peerId);
            WebSocketSession removedSession = PlusWebSocketHandler.peerIdSessionMap.remove(peerId);
            try {
                if (Objects.nonNull(removedSession)) {
                    removedSession.close();
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            } finally {
                PlusWebSocketHandler.peerIdSessionMap.put(peerId, session);
                refreshOnlineSessionsList();
            }
            return;
        }
        //ai翻译
        LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);

        if (ObjectUtil.isNull(loginUser)) {
            session.close(CloseStatus.BAD_DATA);
            log.info("[connect] invalid token received. sessionId: {}", session.getId());
            return;
        }
        WebSocketSessionHolder.addSession(loginUser.getUserId(), session);
        log.info("[connect] sessionId: {},userId:{},userType:{}", session.getId(), loginUser.getUserId(), loginUser.getUserType());
    }

    /**
     * 处理接收到的文本消息
     *
     * @param session WebSocket会话
     * @param message 接收到的文本消息
     * @throws Exception 处理消息过程中可能抛出的异常
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        WORKER_POOL.submit(() -> {

            //视频通话AI消息
            JSONObject msg = JSON.parseObject(message.getPayload());
            switch (msg.getString("type")) {
                case "lang":
                    JSONObject langS = msg.getJSONObject("data");
                    languageMap.put(session.getId(), langS.getString("he"));
                    msg.put("code", 200);
                    WebSocketUtils.sendMessage(session, msg.toJSONString());
                    break;
                case "ping":
                    replyDTO ping = new replyDTO("ping", null);
                    WebSocketUtils.sendMessage(session, JSON.toJSONString(ping));
                    break;
                case "video":
                    //保存视频通话id
                    videoAiMap.remove(msg.getString("oppositeId"));
                    videoAiMap.remove(msg.getString("ourId"));
                    //保存视频通话id
                    videoAiMap.put(msg.getString("oppositeId"), msg.getString("ourId"));
                    videoAiMap.put(msg.getString("ourId"), msg.getString("oppositeId"));
                    replyDTO video = new replyDTO("video", null);
                    WebSocketUtils.sendMessage(session, JSON.toJSONString(video));
                    break;
                case "translate":
                    Boolean status = msg.getBoolean("data");
                    if (status) {
                        statusMap.put(session.getId(), true);
                    } else {
                        statusMap.put(session.getId(), false);
                    }
                    replyDTO translate = new replyDTO("translate", null);
                    WebSocketUtils.sendMessage(session, JSON.toJSONString(translate));
                    break;
                case "text":
                    //获取消息
                    String text = msg.getString("data");
                    try {
                        // 从WebSocket会话中获取登录用户信息
//        LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
//        log.info("msg:{}", message.getPayload());
                        String response = text;
                        aiChatService.saveMessage(response, null, LlmSender.USER, null);
                        //获取语言
                        String lang = languageMap.getOrDefault(session.getId(), "英语");
                        String aiResponse = aiTranslateService.translate(response, lang);
                        log.info("大模型返回信息: {}，回复给用户", aiResponse);
                        String uuid = UUID.randomUUID().toString().replace("-", "");
                        JSONObject texts = new JSONObject();
                        texts.put("user", response);
                        texts.put("ai", aiResponse);
                        texts.put("id", uuid);
                        replyDTO replyText = new replyDTO("text", texts);
                        WebSocketUtils.sendMessage(session, JSON.toJSONString(replyText));
                        String audioPath = aiChatService.generateAudio(aiResponse);
//                log.info("语音转换完成，准备回复给用户语音byte数组: {}", audioPath);
                        byte[] bytes = UUIDToByteArray.uuidToBytes(uuid);
                        byte[] audioData = AudioUtils.mp3ToPcm(audioPath);
                        // 使用ByteBuffer简化合并操作
                        byte[] combined = ByteBuffer.allocate(bytes.length + audioData.length)
                            .put(bytes)
                            .put(audioData)
                            .array();
                        if (session != null && session.isOpen()) {
                            session.sendMessage(new BinaryMessage(combined));
                        }
                        //保存聊天记录
                        aiChatService.saveMessage(aiResponse, null, LlmSender.ASSISTANT, audioPath);
                    } catch (Exception e) {
                        log.error("处理音频异常:{}", e.getMessage());
                    }
                    break;
            }
        });
    }

    /**
     * 处理接收到的二进制消息
     *
     * @param session WebSocket会话
     * @param message 接收到的二进制消息
     * @throws Exception 处理消息过程中可能抛出的异常
     */
    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        //视频通话AI消息
        String peerId = (String) session.getAttributes().get(PEER_ID);
        if (peerId != null) {
            WORKER_POOL.submit(() -> {
                // 发送二进制数据
                try {
                    //回复
                    super.handleBinaryMessage(session, message);
                    log.info("可以在这里进行AI处理，调用大模型接口");
                    // 将耗时操作提交到线程池处理，避免阻塞IO线程
                    String response = aiChatService.audioRecognition(session.getId(), message.getPayload().array());
                    // 判断是否为空或空白字符串
                    if (StringUtils.isBlank(response)) {
                        log.info("没有识别到语音或返回无效内容");
                        return;
                    }
                    //
                    log.info("ASR返回信息: {}，用户说话信息", response);
                    aiChatService.saveMessage(response, null, LlmSender.USER, null);
                    //获取语言
                    String lang = languageMap.getOrDefault(session.getId(), "英语");
                    String aiResponse = aiTranslateService.translate(response, lang);
                    log.info("大模型返回信息: {}，回复给用户", aiResponse);
                    //将消息发送给自己
                    JSONObject jsonChat = new JSONObject();
                    jsonChat.put("type", "my");
                    jsonChat.put("user", response);
                    jsonChat.put("ai", aiResponse);
                    replyDTO reply = new replyDTO("reply", jsonChat);
                    WebSocketUtils.sendMessage(session, JSON.toJSONString(reply));
                    //将消息发送给对方
                    String id = videoAiMap.get(peerId);
                    WebSocketSession webSocketSession = peerIdSessionMap.get(id);
                    jsonChat.put("type", "he");
                    reply.setData(jsonChat);
                    WebSocketUtils.sendMessage(webSocketSession, JSON.toJSONString(reply));
                    //保存聊天记录
                    aiChatService.saveMessage(aiResponse, null, LlmSender.ASSISTANT, null);
                } catch (Exception e) {
                    log.error("处理音频异常:{}", e.getMessage());
                }
            });
            return;
        }
        Boolean put = statusMap.get(session.getId());
        if (put != null && put) {
            //处理普通音频数据
            WORKER_POOL.submit(() -> {
                // 发送二进制数据
                try {
                    //回复
                    super.handleBinaryMessage(session, message);
                    log.info("可以在这里进行AI处理，调用大模型接口");
                    // 将耗时操作提交到线程池处理，避免阻塞IO线程
                    String response = aiChatService.audioRecognition(session.getId(), message.getPayload().array());
                    // 判断是否为空或空白字符串
                    if (StringUtils.isBlank(response)) {
                        log.info("没有识别到语音或返回无效内容");
                        return;
                    }
                    log.info("ASR返回信息: {}，用户说话信息", response);
                    String aiResponse = aiTranslateService.revise(response);
                    aiChatService.saveMessage(aiResponse, null, LlmSender.USER, null);
                    String uuid = UUID.randomUUID().toString().replace("-", "");
                    JSONObject jsonChat = new JSONObject();
                    jsonChat.put("user", aiResponse);
                    jsonChat.put("id", uuid);
                    replyDTO reply = new replyDTO("reply", jsonChat);
                    WebSocketUtils.sendMessage(session, JSON.toJSONString(reply));
                    //保存聊天记录
                    aiChatService.saveMessage(null, null, LlmSender.ASSISTANT, null);
                } catch (Exception e) {
                    log.error("处理音频异常:{}", e.getMessage());
                }
            });
            return;
        }
        //处理普通音频数据
        WORKER_POOL.submit(() -> {
            // 发送二进制数据
            try {
                //回复
                super.handleBinaryMessage(session, message);
                log.info("可以在这里进行AI处理，调用大模型接口");
                // 将耗时操作提交到线程池处理，避免阻塞IO线程
                String response = aiChatService.audioRecognition(session.getId(), message.getPayload().array());
                // 判断是否为空或空白字符串
                if (StringUtils.isBlank(response)) {
                    log.info("没有识别到语音或返回无效内容");
                    return;
                }
                //
                log.info("ASR返回信息: {}，用户说话信息", response);
                aiChatService.saveMessage(response, null, LlmSender.USER, null);
                //获取语言
                String lang = languageMap.getOrDefault(session.getId(), "英语");
                String aiResponse = aiTranslateService.translate(response, lang);
                log.info("大模型返回信息: {}，回复给用户", aiResponse);
                String uuid = UUID.randomUUID().toString().replace("-", "");
                JSONObject jsonChat = new JSONObject();
                jsonChat.put("user", response);
                jsonChat.put("ai", aiResponse);
                jsonChat.put("id", uuid);
                replyDTO reply = new replyDTO("reply", jsonChat);
                WebSocketUtils.sendMessage(session, JSON.toJSONString(reply));
                String audioPath = aiChatService.generateAudio(aiResponse);
//                log.info("语音转换完成，准备回复给用户语音byte数组: {}", audioPath);
                // 发送二进制数据
                byte[] bytes = UUIDToByteArray.uuidToBytes(uuid);
                byte[] audioData = AudioUtils.mp3ToPcm(audioPath);
                // 使用ByteBuffer简化合并操作
                byte[] combined = ByteBuffer.allocate(bytes.length + audioData.length)
                    .put(bytes)
                    .put(audioData)
                    .array();
                if (session != null && session.isOpen()) {
                    session.sendMessage(new BinaryMessage(combined));
                }
                //保存聊天记录
                aiChatService.saveMessage(aiResponse, null, LlmSender.ASSISTANT, audioPath);
            } catch (Exception e) {
                log.error("处理音频异常:{}", e.getMessage());
            }
        });

    }

    /**
     * 处理接收到的Pong消息（心跳监测）
     *
     * @param session WebSocket会话
     * @param message 接收到的Pong消息
     * @throws Exception 处理消息过程中可能抛出的异常
     */
    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        WebSocketUtils.sendPongMessage(session);
    }

    /**
     * 处理WebSocket传输错误
     *
     * @param session   WebSocket会话
     * @param exception 发生的异常
     * @throws Exception 处理过程中可能抛出的异常
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("[transport error] sessionId: {} , exception:{}", session.getId(), exception.getMessage());
        languageMap.remove(session.getId());

        //视频通话消息
        String peerId = (String) session.getAttributes().get(PEER_ID);
        if (peerId != null) {
            WebSocketSession removedSession = PlusWebSocketHandler.peerIdSessionMap.remove(peerId);
            try {
                if (Objects.nonNull(removedSession)) {
                    removedSession.close();
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            } finally {
                refreshOnlineSessionsList();
            }
        }
    }

    /**
     * 在WebSocket连接关闭后执行清理操作
     *
     * @param session WebSocket会话
     * @param status  关闭状态信息
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        languageMap.remove(session.getId());
        statusMap.remove(session.getId());
        //视频通话消息
        String peerId = (String) session.getAttributes().get(PEER_ID);
        if (peerId != null) {
            log.warn("on close:the session is is :{},the peer id is:{}", session.getId(), peerId);
            WebSocketSession removedSession = PlusWebSocketHandler.peerIdSessionMap.remove(peerId);
            try {
                if (Objects.nonNull(removedSession)) {
                    removedSession.close();
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            } finally {
                refreshOnlineSessionsList();
            }
            return;
        }
        //普通对话消息
        LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
        if (ObjectUtil.isNull(loginUser)) {
            log.info("[disconnect] invalid token received. sessionId: {}", session.getId());
            return;
        }
        WebSocketSessionHolder.removeSession(loginUser.getUserId());
        log.info("[disconnect] sessionId: {},userId:{},userType:{}", session.getId(), loginUser.getUserId(), loginUser.getUserType());
    }

    /**
     * 指示处理程序是否支持接收部分消息
     *
     * @return 如果支持接收部分消息，则返回true；否则返回false
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 视频通话上线人员列表
     */
    private void refreshOnlineSessionsList() {
        PlusWebSocketHandler.peerIdSessionMap.forEach((peerId, session) -> {
            JSONObject onlineSession = new JSONObject();
            onlineSession.put("type", "online");
            onlineSession.put("data", JSON.toJSONString(PlusWebSocketHandler.peerIdSessionMap.keySet()));
            WebSocketUtils.sendMessage(session, onlineSession.toJSONString());
        });
    }

}
